{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a30ea8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b752f5a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mesolitica/.local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from accelerate import init_empty_weights, load_checkpoint_and_dispatch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer\n",
    "import safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "982b9e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-Small-3.1-24B-Instruct-2503')\n",
    "x = {\n",
    "  \"chat_template\": \"{%- set today = strftime_now(\\\"%Y-%m-%d\\\") %}\\n{%- set default_system_message = \\\"You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\\\\nYour knowledge base was last updated on 2023-10-01. The current date is \\\" + today + \\\".\\\\n\\\\nWhen you're not sure about some information, you say that you don't have the information and don't make up anything.\\\\nIf the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. \\\\\\\"What are some good restaurants around me?\\\\\\\" => \\\\\\\"Where are you?\\\\\\\" or \\\\\\\"When is the next flight to Tokyo\\\\\\\" => \\\\\\\"Where do you travel from?\\\\\\\")\\\" %}\\n\\n{{- bos_token }}\\n\\n{%- if messages[0]['role'] == 'system' %}\\n    {%- if messages[0]['content'] is string %}\\n        {%- set system_message = messages[0]['content'] %}\\n    {%- else %}\\n        {%- set system_message = messages[0]['content'][0]['text'] %}\\n    {%- endif %}\\n    {%- set loop_messages = messages[1:] %}\\n{%- else %}\\n    {%- set system_message = default_system_message %}\\n    {%- set loop_messages = messages %}\\n{%- endif %}\\n{{- '[SYSTEM_PROMPT]' + system_message + '[/SYSTEM_PROMPT]' }}\\n\\n{%- for message in loop_messages %}\\n    {%- if message['role'] == 'user' %}\\n        {%- if message['content'] is string %}\\n            {{- '[INST]' + message['content'] + '[/INST]' }}\\n        {%- else %}\\n            {{- '[INST]' }}\\n            {%- for block in message['content'] %}\\n                {%- if block['type'] == 'text' %}\\n                    {{- block['text'] }}\\n                {%- elif block['type'] in ['image', 'image_url'] %}\\n                    {{- '[IMG]' }}\\n                {%- else %}\\n                    {{- raise_exception('Only text and image blocks are supported in message content!') }}\\n                {%- endif %}\\n            {%- endfor %}\\n            {{- '[/INST]' }}\\n        {%- endif %}\\n    {%- elif message['role'] == 'system' %}\\n        {%- if message['content'] is string %}\\n            {{- '[SYSTEM_PROMPT]' + message['content'] + '[/SYSTEM_PROMPT]' }}\\n        {%- else %}\\n            {{- '[SYSTEM_PROMPT]' + message['content'][0]['text'] + '[/SYSTEM_PROMPT]' }}\\n        {%- endif %}\\n    {%- elif message['role'] == 'assistant' %}\\n        {%- if message['content'] is string %}\\n            {{- message['content'] + eos_token }}\\n        {%- else %}\\n            {{- message['content'][0]['text'] + eos_token }}\\n        {%- endif %}\\n    {%- else %}\\n        {{- raise_exception('Only user, system and assistant roles are supported!') }}\\n    {%- endif %}\\n{%- endfor %}\"\n",
    "}\n",
    "tokenizer.chat_template = x['chat_template']\n",
    "streamer = TextStreamer(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5d3f19b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-06-01 16:47:16,704] [INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/bin/ld: cannot find -laio: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n",
      "/usr/bin/ld: cannot find -lcufile: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n",
      "Loading checkpoint shards: 100%|████████████████████████████████| 10/10 [00:09<00:00,  1.11it/s]\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    './local', torch_dtype = 'auto',\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76afcba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict = model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b9c6e367",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoint-2900  checkpoint-3000  checkpoint-3100\r\n"
     ]
    }
   ],
   "source": [
    "!ls lora-embedding-256-Mistral-Small-3.1-24B-Instruct-2503-malaysian-8k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7dd07b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "from safetensors import safe_open\n",
    "\n",
    "folder = 'lora-embedding-256-Mistral-Small-3.1-24B-Instruct-2503-malaysian-8k/checkpoint-3100'\n",
    "\n",
    "f = safe_open(f\"{folder}/adapter_model.safetensors\", framework=\"pt\", device='cpu')\n",
    "keys = f.keys()\n",
    "keys = sorted(list(set([k.split('.lora')[0] for k in keys if '.lora' in k])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eba28fdb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████| 282/282 [00:00<00:00, 379.12it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "for k in tqdm(keys):\n",
    "    k_ori = k.replace('base_model.model.', '') + '.weight'\n",
    "    if 'embed_tokens' in k:\n",
    "        post_A = '.lora_embedding_A'\n",
    "        post_B = '.lora_embedding_B'\n",
    "    else:\n",
    "        post_A = '.lora_A.weight'\n",
    "        post_B = '.lora_B.weight'\n",
    "    A = k + post_A\n",
    "    B = k + post_B\n",
    "    \n",
    "    W = state_dict[k_ori]\n",
    "    if 'embed_tokens' not in k:\n",
    "        W = W.t()\n",
    "        \n",
    "    A = f.get_tensor(A).to(W.device)\n",
    "    B = f.get_tensor(B).to(W.device)\n",
    "    with torch.no_grad():\n",
    "        W.addmm_(A.t(), B.t(), alpha = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cd50b3fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s>[SYSTEM_PROMPT]You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n",
      "Your knowledge base was last updated on 2023-10-01. The current date is 2025-06-01.\n",
      "\n",
      "When you're not sure about some information, you say that you don't have the information and don't make up anything.\n",
      "If the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. \"What are some good restaurants around me?\" => \"Where are you?\" or \"When is the next flight to Tokyo\" => \"Where do you travel from?\")[/SYSTEM_PROMPT][INST]\"butoh hang\" tu apa dalam bahsa kedah[/INST]Ungkapan yang digunakan untuk menggambarkan seseorang atau sesuatu yang kelihatan sangat buruk, tidak menyenangkan, atau menjijikkan.. \"Butoh\" bermaksud hancur atau rosak teruk dalam bahasa Kedah, manakala \"hang\" adalah kata ganti nama diri ketiga untuk lelaki. Gabungan kedua-dua perkataan ini menghasilkan ungkapan yang kuat dan ekspresif untuk menggambarkan penampilan yang amat tidak memuaskan.</s>\n"
     ]
    }
   ],
   "source": [
    "d = [\n",
    "    {'role': 'user', 'content': '\"butoh hang\" tu apa dalam bahsa kedah'}\n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(d, return_tensors = 'pt', add_generation_prompt=True).to('cuda')\n",
    "generate_kwargs = dict(\n",
    "    input_ids=inputs,\n",
    "    max_new_tokens=128,\n",
    "    top_p=0.95,\n",
    "    top_k=50,\n",
    "    temperature=0.9,\n",
    "    do_sample=True,\n",
    "    repetition_penalty=1.05,\n",
    "    streamer=streamer\n",
    ")\n",
    "generation_output = model.generate(**generate_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "117e2db4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s>[SYSTEM_PROMPT]You are going to enter reasoning mode, first you try to think step-by-step in malay after that give the final answer.[/SYSTEM_PROMPT][INST]\"Awok ni kene bebele dengan dokte, bomo tok jalan.\" terjemah ke malay[/INST]Baguslah. Kita mulakan proses penterjemahan ayat dialek Pahang ini secara **panjang lebar, terperinci dan bernaratif**, seperti yang anda minta. Ini bukan sekadar menggantikan perkataan demi perkataan kepada Bahasa Melayu Standard, tetapi juga memerlukan pemahaman konteks budaya, nuansa emosi, dan makna tersirat dalam ayat tersebut.\n",
      "\n",
      "## **Langkah 1: Memahami Fungsi dan Nada Ayat Keseluruhan**\n",
      "\n",
      "Ayat ini kelihatan seperti satu nasihat atau teguran ringan, mungkin dalam situasi di mana seseorang berhadapan dengan cabaran fizikal atau halangan, dan disarankan untuk menggunakan tenaga dan strategi yang betul supaya dapat mengatasi masalah tersebut.\n",
      "\n",
      "Nada pertuturan adalah **bersahabat dan memberi sokongan** – penutur memberikan dorongan positif dan keyakinan bahawa jika usaha dilakukan dengan betul, maka kejayaan boleh dicapai.\n",
      "\n",
      "> Fungsi ayat: Nasihat dan dorongan.\n",
      "> Nada: Positif, menyokong dan memberi semangat.\n",
      "\n",
      "## **Langkah 2: Pecahan Ayat kepada Unit Tatabahasa Kecil**\n",
      "\n",
      "Mari kita pecahkan ayat ini kepada enam bahagian utama:\n",
      "\n",
      "1. **\"Awok ni\"**\n",
      "2. **\"Kene bejela\"**\n",
      "3. **\"Dengan dokte\"**\n",
      "4. **\"Bomo tok jalan\"**\n",
      "\n",
      "Setiap unit ini akan dianalisis secara berasingan untuk memahami peranan dan maknanya dalam keseluruhan ayat.\n",
      "\n",
      "## **Langkah 3: Analisis Semantik dan Sintaksis Setiap Perkataan**\n",
      "\n",
      "### 1. **\"Awok ni\"**\n",
      "\n",
      "* **\"Awok\"** dalam dialek Pahang merujuk kepada \"awak\" atau \"anda\" – kata ganti nama diri kedua.\n",
      "* **\"Ni\"** adalah bentuk pendek bagi \"ini\", tetapi dalam konteks ini, ia berfungsi sebagai **penegas** atau penekanan pada subjek.\n",
      "\n",
      "### 2. **\"Kene bejela\"**\n",
      "\n",
      "* **\"Kene\"** bermaksud \"harus\" atau \"perlu\" – bentuk imperatif atau syarat.\n",
      "* **\"Bejela\"** ialah perkataan yang tidak wujud dalam kamus standard, tetapi berdasarkan konteks dan analogi dengan dialek lain, ia mungkin merujuk kepada \"berlari\" atau \"berhenti\". Dalam dialek Negeri Sembilan, \"belajau\" bermaksud \"lari\", dan mungkin ada hubungan.\n",
      "* Jadi, \"kene bejela\" boleh dianggap sebagai \"harus berlari\".\n",
      "\n",
      "### 3. **\"Dengan dokte\"**\n",
      "\n",
      "* **\"Dengan\"** digunakan untuk menghubungkan tindakan dengan sumber daya atau alat.\n",
      "* **\"Dokte\"** mungkin merujuk kepada \"dokter\" (juru bedah), tetapi dalam konteks ini, lebih masuk akal jika ia merujuk kepada \"duit\" (wang) kerana ia sering kali diperlukan untuk melakukan banyak perkara.\n",
      "\n",
      "### 4. **\"Bomo tok jalan\"**\n",
      "\n",
      "* **\"Bomo\"** mungkin merujuk kepada \"bomoh\" atau bomoh – individu yang mempunyai kemahiran unik atau \"trick\".\n",
      "* **\"Tok jalan\"** boleh dilihat sebagai bentuk dialek bagi \"tak jalan\" atau \"tidak jalan\" – iaitu tiada jalan atau cara.\n",
      "\n",
      "## **Langkah 4: Penyusunan Maksud Secara Literal Berdasarkan Komponen**\n",
      "\n",
      "Berdasarkan analisis komponen-komponen tadi, kita cuba menyusun semula maksud literal secara bertahap:\n",
      "\n",
      "1. **\"Awok ni\"** -> Awak ini\n",
      "2. **\"Kene bejela\"** -> Harus berlari\n",
      "3. **\"Dengan dokte\"** -> Dengan duit\n",
      "4. **\"Bomo tok jalan\"** -> Bomoh tak jalan\n",
      "\n",
      "## **Langkah 5: Transformasi ke Bentuk Bahasa Melayu Standard**\n",
      "\n",
      "Kita kini akan menyusun semula ayat ini menjadi **beberapa versi** dalam Bahasa Melayu Standard, dengan pelbagai tahap keformalan dan kehalusan bahasa:\n",
      "\n",
      "### **Versi A – Langsung dan Ringkas**\n",
      "\n",
      "> \"Awak ni harus berlari dengan duit, bomoh tak jalan.\"\n",
      "\n",
      "### **Versi B – Sopan dan Naratif**\n",
      "\n",
      "> \"Anda perlu berusaha keras dengan sumber daya yang ada, kerana tiada cara lain selain itu.\"\n",
      "\n",
      "### **Versi C – Formal dan Puitis**\n",
      "\n",
      "> \"Untuk mencapai matlamat, usahawan itu harus bergerak pantas dan cekap dengan modal yang dimilikinya, kerana tiada alternatif yang lebih baik.\"\n",
      "\n",
      "## **Langkah 6: Analisis Makna Tersirat dan Konteks Emosi**\n",
      "\n",
      "Dalam konteks budaya Pahang yang sering penuh dengan nilai ketekunan dan kepintaran, ayat ini boleh difahami sebagai **nasihat praktikal** tentang bagaimana menangani masalah atau halangan. Ia menekankan **pentingnya usaha sendiri** dan penggunaan sumber daya yang sesuai, tanpa bergantung kepada faktor luaran atau \"bomoh\".\n",
      "\n",
      "## **Kesimpulan Akhir**\n",
      "\n",
      "Ayat dialek Pahang ini menunjukkan bagaimana bahasa tempatan dapat mengungkapkan nasihat praktikal dengan cara yang pendek dan padat. Penterjemahan yang paling tepat dan padat ialah:\n",
      "\n",
      "> \"Engkau harus berlari dengan duit, bomoh tak jalan.\"\n",
      "\n",
      "Ayat ini bukan sahaja menyampaikan mesej literal, tetapi juga membawa **nilai moral dan etika** mengenai kerja keras dan kecerdikan.\n",
      "\n",
      "**Jawapan Akhir**\n",
      "\n",
      "$boxed{Engkau harus berlari dengan duit, bomoh tak jalan.}$.</s>\n"
     ]
    }
   ],
   "source": [
    "d = [\n",
    "    {'role': 'system', 'content': 'You are going to enter reasoning mode, first you try to think step-by-step in malay after that give the final answer.'},\n",
    "    {'role': 'user', 'content': '\"Awok ni kene bebele dengan dokte, bomo tok jalan.\" terjemah ke malay'},\n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(d, return_tensors = 'pt').to('cuda')\n",
    "generate_kwargs = dict(\n",
    "    input_ids=inputs,\n",
    "    max_new_tokens=4096,\n",
    "    top_p=0.95,\n",
    "    top_k=50,\n",
    "    temperature=0.9,\n",
    "    do_sample=True,\n",
    "    repetition_penalty=1.05,\n",
    "    streamer=streamer\n",
    ")\n",
    "generation_output = model.generate(**generate_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "52be6e8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503/commit/04e2a5cc3e291ba13baf94f21e337847d46f3613', commit_message='Upload tokenizer', commit_description='', oid='04e2a5cc3e291ba13baf94f21e337847d46f3613', pr_url=None, repo_url=RepoUrl('https://huggingface.co/mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503', endpoint='https://huggingface.co', repo_type='model', repo_id='mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "31917105",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('./24b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a9eb052",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start hashing 13 files.\n",
      "Finished hashing 13 files.\n",
      "Uploading files using Xet Storage..\n",
      "Uploading...: 100%|████████████████████████| 47.1G/47.1G [10:39<00:00, 73.7MB/s]\n",
      "https://huggingface.co/mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503/tree/main/.\n"
     ]
    }
   ],
   "source": [
    "!huggingface-cli upload mesolitica/Malaysian-Mistral-Small-3.1-24B-Instruct-2503 ./24b ."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
